# cython: language_level=3
# -*- coding: utf-8 -*-
"""
轨迹基础类模块
提供轨迹数据存储和管理的基础功能

该模块定义了TRAJ_BASE类，用于存储和管理强化学习中的轨迹数据，
包括状态、动作、奖励等时序数据的存储、类型检查和内存管理。
"""

import numpy as np


class TRAJ_BASE():
    """
    轨迹基础类

    这是一个用于存储强化学习轨迹数据的基础类。它提供了动态添加数据项、
    类型检查、内存管理等功能。每个轨迹实例可以存储多个不同类型的数据序列。

    Attributes:
        key_data_type (dict): 存储各个数据键对应的数据类型，类变量
        key_data_shape (dict): 存储各个数据键对应的数据形状，类变量
        max_mem_length (int): 记录所有轨迹实例中的最大内存长度，类变量
    """

    # 类变量：用于记录所有轨迹实例中各个数据键的类型和形状信息
    key_data_type = {}  # 数据类型字典 {key: dtype}
    key_data_shape = {}  # 数据形状字典 {key: shape}
    max_mem_length = -1  # 最大内存长度，用于内存优化监控

    def __init__(self, traj_limit, env_id):
        """
        初始化轨迹基础对象

        Args:
            traj_limit (int): 轨迹的最大长度限制（时间步数）
            env_id: 环境标识符，用于区分不同的环境
        """
        self.traj_limit = traj_limit  # 轨迹长度限制
        self.env_id = env_id  # 环境ID
        self.readonly_lock = False  # 只读锁，锁定后不能再修改数据
        self.key_dict = []  # 存储所有数据键的列表
        self.time_pointer = 0  # 时间指针，指向当前要写入数据的位置
        self.need_reward_bootstrap = False  # 是否需要奖励回补标志
        self.deprecated_flag = False  # 废弃标志，标记轨迹是否已废弃

    def set_scalar(self, key, value):
        """
        设置标量属性

        为对象添加一个标量属性，用于存储非序列化的数据。

        Args:
            key (str): 属性名
            value: 属性值
        """
        setattr(self, key, value)

    def remember(self, key, content):
        """
        记忆数据：在当前时间步存储数据

        这是核心的数据存储方法，用于在轨迹中记录某个时间步的数据。
        支持动态添加新的数据键，自动处理数据类型和形状检查。

        Args:
            key (str): 数据键名，如 'state', 'action', 'reward' 等
            content: 要存储的数据内容，可以是numpy数组或标量，
                    None表示该时间步没有对应数据

        处理逻辑：
        1. 如果是新的数据键且有内容 -> 初始化数据结构并存储内容
        2. 如果是新的数据键但无内容 -> 初始化空的数据结构
        3. 如果是已存在的数据键且有内容 -> 直接存储内容
        4. 如果是已存在的数据键但无内容 -> 不做任何操作
        """
        assert not self.readonly_lock, "轨迹已锁定，无法修改数据"

        # 情况1：新数据键且有内容 - 需要初始化数据结构
        if not (key in self.key_dict) and (content is not None):
            self.init_track(key=key, first_content=content)  # 初始化数据结构
            getattr(self, key)[self.time_pointer] = content   # 存储当前时间步的数据

        # 情况2：新数据键但无内容 - 初始化空的数据结构
        elif not (key in self.key_dict) and (content is None):
            self.init_track_none(key=key)  # 初始化空的数据结构

        # 情况3：已存在的数据键且有内容 - 直接存储
        elif (key in self.key_dict) and (content is not None):
            getattr(self, key)[self.time_pointer] = content

        # 情况4：已存在的数据键但无内容 - 不做操作
        else:
            pass

    def copy_track(self, origin_key, new_key):
        """
        复制/重命名轨迹数据

        将一个已存在的数据键复制到新的数据键，支持嵌套键（使用'>'分隔）。
        主要用于数据预处理或特征工程时创建数据的副本。

        Args:
            origin_key (str): 原始数据键名
            new_key (str): 新的数据键名

        处理逻辑：
        1. 如果直接匹配到原始键 -> 简单复制数据
        2. 如果没有直接匹配但找到嵌套键 -> 批量复制所有匹配的嵌套键
        """
        # 情况1：直接匹配到原始键
        if hasattr(self, origin_key):
            origin_handle = getattr(self, origin_key)  # 获取原始数据
            setattr(self, new_key, origin_handle.copy())  # 创建副本
            new_handle = getattr(self, new_key)
            self.key_dict.append(new_key)  # 添加新键到键字典

        # 情况2：处理嵌套键（例如：'obs>image' 这样的键）
        else:
            # 查找所有包含原始键作为前缀的嵌套键
            real_key_list = [real_key for real_key in self.__dict__
                           if (origin_key + '>' in real_key)]
            assert len(real_key_list) > 0, f"键不存在: {origin_key}"

            # 批量复制所有匹配的嵌套键
            for real_key in real_key_list:
                mainkey, subkey = real_key.split('>')  # 分割主键和子键
                self.copy_track(real_key, (new_key + '>' + subkey))

    def check_type_shape(self, key, first_content=None):
        """
        检查并确保数据类型和形状的一致性

        这是类型安全管理的关键方法，确保同一数据键在不同时间步的数据
        具有相同的类型和形状。支持类型注册和一致性验证。

        Args:
            key (str): 数据键名
            first_content: 首次出现的数据内容，用于注册类型信息

        Returns:
            tuple: (数据类型, 数据形状)

        功能：
        1. 如果提供首次内容 -> 注册或验证类型信息
        2. 如果不提供内容 -> 从已注册信息中获取类型
        """
        # 情况1：提供首次内容，需要注册或验证类型
        if first_content is not None:
            content_type = first_content.dtype  # 获取数据类型
            content_shape = first_content.shape  # 获取数据形状

            # 检查类型一致性
            if key in TRAJ_BASE.key_data_type:
                assert TRAJ_BASE.key_data_type[key] == content_type, \
                    f"数据类型不匹配: {key}"
            else:
                # 首次出现，注册类型信息
                TRAJ_BASE.key_data_type[key] = content_type
                TRAJ_BASE.key_data_shape[key] = content_shape

            return content_type, content_shape

        # 情况2：从已注册信息获取类型
        assert key in TRAJ_BASE.key_data_type, \
            f"未找到键的类型信息: {key}"
        return TRAJ_BASE.key_data_type[key], TRAJ_BASE.key_data_shape[key]

    def init_track(self, key, first_content):
        """
        初始化数据轨道

        当某个数据键首次出现时调用，创建合适大小的numpy数组来存储该数据。
        根据数据的类型和形状预分配内存，提高后续存储效率。

        Args:
            key (str): 数据键名
            first_content: 首个数据内容，用于确定数据类型和形状
        """
        content = first_content
        self.check_type_shape(key, first_content)  # 检查并注册类型信息

        # 确保内容是numpy数组或标量
        assert isinstance(content, np.ndarray) or isinstance(content, np.ScalarType), \
            f"不支持的数据类型: {key}, {content.__class__}"

        # 计算张量大小：(轨迹长度, 数据形状...)
        tensor_size = ((self.traj_limit,) + tuple(content.shape))

        # 创建预分配的数组
        set_item = np.zeros(shape=tensor_size, dtype=content.dtype)

        # 根据数据类型设置默认值
        # 浮点类型用NaN填充，其他类型用0填充
        set_item[:] = np.nan if np.issubdtype(content.dtype, np.floating) else 0

        # 设置为对象属性并记录键名
        setattr(self, key, set_item)
        self.key_dict.append(key)

    def init_track_none(self, key):
        """
        初始化空数据轨道

        当某个数据键首次出现但内容为None时调用。从历史类型字典中
        读取类型信息来创建合适的数据结构。

        Args:
            key (str): 数据键名
        """
        # 从类变量中获取已注册的类型信息
        content_dtype, content_shape = self.check_type_shape(key)

        # 计算张量大小
        tensor_size = ((self.traj_limit,) + tuple(content_shape))

        # 创建预分配数组
        set_item = np.zeros(shape=tensor_size, dtype=content_dtype)

        # 设置默认值
        set_item[:] = np.nan if np.issubdtype(content_dtype, np.floating) else 0

        # 设置属性并记录键名
        setattr(self, key, set_item)
        self.key_dict.append(key)

    def time_shift(self):
        """
        时间指针推进

        将时间指针向前推进一位，准备记录下一个时间步的数据。
        在每次完成当前时间步的数据记录后调用此方法。

        Raises:
            AssertionError: 如果时间指针已达到轨迹长度限制
        """
        assert self.time_pointer < self.traj_limit, \
            f"时间指针超出轨迹限制: {self.time_pointer} >= {self.traj_limit}"
        self.time_pointer += 1

    def cut_tail(self):
        """
        截断轨迹尾部

        当轨迹的实际长度小于预分配的最大长度时，截断未使用的尾部空间。
        这在轨迹完成时调用，可以节省内存使用。

        功能：
        1. 锁定轨迹为只读状态
        2. 更新最大内存长度记录
        3. 截断所有数据数组的尾部未使用空间
        """
        # 辅助函数：获取指定键的数据
        TJ = lambda key: getattr(self, key)

        # 锁定轨迹，防止后续修改
        self.readonly_lock = True

        # 获取实际使用的帧数
        n_frame = self.time_pointer

        # 检查并更新最大内存长度记录（用于内存监控）
        if n_frame > TRAJ_BASE.max_mem_length:
            TRAJ_BASE.max_mem_length = n_frame
            print(f'内存使用更新 - 最大长度: {TRAJ_BASE.max_mem_length}, 轨迹限制: {self.traj_limit}')

        # 截断所有数据数组的尾部空间
        for key in self.key_dict:
            setattr(self, key, TJ(key)[:n_frame])